package com.me.common.util.security;

import org.bouncycastle.bcpg.*;
import org.bouncycastle.bcpg.sig.Features;
import org.bouncycastle.bcpg.sig.KeyFlags;
import org.bouncycastle.crypto.generators.RSAKeyPairGenerator;
import org.bouncycastle.crypto.params.RSAKeyGenerationParameters;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.openpgp.*;
import org.bouncycastle.openpgp.jcajce.JcaPGPObjectFactory;
import org.bouncycastle.openpgp.operator.PBESecretKeyEncryptor;
import org.bouncycastle.openpgp.operator.PGPDigestCalculator;
import org.bouncycastle.openpgp.operator.bc.*;
import org.bouncycastle.openpgp.operator.jcajce.JcaKeyFingerprintCalculator;
import org.bouncycastle.openpgp.operator.jcajce.JcaPGPContentSignerBuilder;
import org.bouncycastle.openpgp.operator.jcajce.JcaPGPKeyPair;
import org.bouncycastle.openpgp.operator.jcajce.JcePBESecretKeyDecryptorBuilder;
import org.bouncycastle.util.io.Streams;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.*;
import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.SecureRandom;
import java.security.Security;
import java.sql.Timestamp;
import java.time.LocalDateTime;
import java.time.ZoneId;
import java.util.Base64;
import java.util.Date;
import java.util.Iterator;
import java.util.Objects;

public class PGPUtils {

    private static final Logger LOG = LoggerFactory.getLogger(PGPUtils.class);

    public static void main(String[] args) throws Exception {
        String password = "12345678";
        String priFileName = "pgp-pri.asc";
        String pubFileName = "pgp-pub.asc";
        // 生成密钥对
        // generatePGPKeys(new KeyId("finance.17usoft.com"), password, priFileName, pubFileName);
        generatePGP("finance.17usoft.com", password);

        PGPPublicKey publicKey = readPublicKey("cert/pgp" + File.separator + pubFileName);
        PGPSecretKey secretKey = readSecretKey("cert/pgp" + File.separator + priFileName);
        PGPPrivateKey privateKey = findPrivateKeyFromSecretKey(secretKey, password);
        String msg = "Hello, PGP!";
        // 加密并签名
        String encryptedStr = encryptAndSign(msg, secretKey, privateKey, publicKey);
        // 解密并验签
        String decryptStr = decryptAndVerify(encryptedStr, privateKey, publicKey);
        if (!msg.equals(decryptStr)) {
            throw new RuntimeException("解密失败！");
        }
        System.out.println("解密消息：" + decryptStr);
    }

    /**
     *
     * @param userId userId
     * @param password 私钥密码
     */
    public static void generatePGP(String userId, String password) {
        Security.addProvider(new BouncyCastleProvider());

        try {
            Path path = Paths.get(Paths.get("").toAbsolutePath() + "/testJDK8/src/main/resources/cert/pgp");
            if (!Files.exists(path)) {
                Files.createDirectory(path);
            }

            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA", "BC");
            keyPairGenerator.initialize(2048, new SecureRandom());
            KeyPair keyPair = keyPairGenerator.generateKeyPair();
            PGPKeyPair pgpKeyPair = new JcaPGPKeyPair(PGPPublicKey.RSA_GENERAL, keyPair, new Date());

            PGPSignatureSubpacketGenerator subPacketGenerator = new PGPSignatureSubpacketGenerator();
            subPacketGenerator.setKeyFlags(false, KeyFlags.SIGN_DATA | KeyFlags.CERTIFY_OTHER | KeyFlags.ENCRYPT_COMMS | KeyFlags.ENCRYPT_STORAGE);
            // 过期时间，秒
            subPacketGenerator.setKeyExpirationTime(true, 365 * 24 * 60 * 60L);
            PGPSignatureSubpacketVector subpacketVector = subPacketGenerator.generate();

            PGPSecretKey pgpSecretKey = new PGPSecretKey(
                    PGPSignature.DEFAULT_CERTIFICATION,
                    pgpKeyPair,
                    userId,
                    new BcPGPDigestCalculatorProvider().get(HashAlgorithmTags.SHA256),
                    subpacketVector,
                    null,
                    new BcPGPContentSignerBuilder(pgpKeyPair.getPublicKey().getAlgorithm(), HashAlgorithmTags.SHA256),
                    new BcPBESecretKeyEncryptorBuilder(PGPEncryptedData.CAST5, new BcPGPDigestCalculatorProvider().get(HashAlgorithmTags.SHA256)).build(password.toCharArray()));

            try (ArmoredOutputStream privateKeyOut = new ArmoredOutputStream(Files.newOutputStream(Paths.get(path.toAbsolutePath() + File.separator + "pgp-pri.asc")));
                 ArmoredOutputStream publicKeyOut = new ArmoredOutputStream(Files.newOutputStream(Paths.get(path.toAbsolutePath() + File.separator + "pgp-pub.asc")))) {
                pgpSecretKey.encode(privateKeyOut);
                pgpSecretKey.getPublicKey().encode(publicKeyOut);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private static PGPKeyRingGenerator generateKeyRingGenerator(KeyId keyId, String password, int s2kCount)
            throws Exception {
        RSAKeyPairGenerator kpg = new RSAKeyPairGenerator();
        kpg.init(new RSAKeyGenerationParameters(BigInteger.valueOf(0x10001), new SecureRandom(), 2048, 12));
        PGPKeyPair rsakp_sign = new BcPGPKeyPair(PGPPublicKey.RSA_GENERAL, kpg.generateKeyPair(), new Date());

        // 永不过期
        PGPSignatureSubpacketGenerator signhashgen = new PGPSignatureSubpacketGenerator();
        signhashgen.setKeyFlags(false, KeyFlags.SIGN_DATA | KeyFlags.ENCRYPT_COMMS | KeyFlags.ENCRYPT_STORAGE);
        signhashgen.setPreferredSymmetricAlgorithms(false, new int[] { SymmetricKeyAlgorithmTags.AES_256,
                SymmetricKeyAlgorithmTags.AES_192, SymmetricKeyAlgorithmTags.AES_128 });
        signhashgen.setPreferredHashAlgorithms(false, new int[] { HashAlgorithmTags.SHA256, HashAlgorithmTags.SHA1,
                HashAlgorithmTags.SHA384, HashAlgorithmTags.SHA512, HashAlgorithmTags.SHA224, });
        signhashgen.setPreferredCompressionAlgorithms(false, new int[] { CompressionAlgorithmTags.ZIP,
                CompressionAlgorithmTags.ZLIB, CompressionAlgorithmTags.BZIP2 });
        signhashgen.setFeature(false, Features.FEATURE_MODIFICATION_DETECTION);
        // 永不过期
        signhashgen.setKeyExpirationTime(false, 0L);

        PGPDigestCalculator sha1Calc = new BcPGPDigestCalculatorProvider().get(HashAlgorithmTags.SHA1);
        PGPDigestCalculator sha256Calc = new BcPGPDigestCalculatorProvider().get(HashAlgorithmTags.SHA256);
        PBESecretKeyEncryptor pske = (new BcPBESecretKeyEncryptorBuilder(PGPEncryptedData.AES_256, sha256Calc,
                s2kCount)).build(password.toCharArray());
        PGPKeyRingGenerator keyRingGen = new PGPKeyRingGenerator(PGPSignature.POSITIVE_CERTIFICATION, rsakp_sign, keyId.get(),
                sha1Calc, signhashgen.generate(), null,
                new BcPGPContentSignerBuilder(rsakp_sign.getPublicKey().getAlgorithm(), HashAlgorithmTags.SHA1), pske);

        return keyRingGen;
    }

    /**
     * 加密并签名    <p></p>
     *
     * @param message       原始消息
     * @param secretKey     安全密钥对
     * @param privateKey    PGP私钥
     * @param publicKey     PGP公钥
     */
    public static String encryptAndSign(String message, PGPSecretKey secretKey, PGPPrivateKey privateKey, PGPPublicKey publicKey) throws Exception {
        byte[] clearData = message.getBytes(StandardCharsets.UTF_8);
        try (InputStream in = new ByteArrayInputStream(clearData);
             ByteArrayOutputStream out = new ByteArrayOutputStream()) {
            // 设置 PGP数据格式
            ArmoredOutputStream armored = new ArmoredOutputStream(out);

            // 加密数据生成器：AES-256
            PGPEncryptedDataGenerator encryptedGenerator = new PGPEncryptedDataGenerator(new BcPGPDataEncryptorBuilder(PGPEncryptedData.AES_256)
                    .setSecureRandom(new SecureRandom()).setWithIntegrityPacket(true));
            encryptedGenerator.addMethod(new BcPublicKeyKeyEncryptionMethodGenerator(publicKey));
            // 包装成 加密格式
            OutputStream encryptedOut = encryptedGenerator.open(armored, new byte[4096]);

            // 压缩数据生成器：ZIP压缩
            PGPCompressedDataGenerator compressedDataGenerator = new PGPCompressedDataGenerator(PGPCompressedData.ZIP);
            // 包装成 加密格式 输出流
            OutputStream compressedOut = compressedDataGenerator.open(encryptedOut, new byte[4096]);

            // 签名数据生成器：
            PGPSignatureGenerator signatureGenerator =
                    new PGPSignatureGenerator(new JcaPGPContentSignerBuilder(secretKey.getPublicKey().getAlgorithm(), PGPUtil.SHA256).setProvider("BC"));
            signatureGenerator.init(PGPSignature.BINARY_DOCUMENT, privateKey);

            Iterator<String> it = secretKey.getPublicKey().getUserIDs();
            if (it.hasNext()) {
                PGPSignatureSubpacketGenerator spGen = new PGPSignatureSubpacketGenerator();
                spGen.setSignerUserID(false, it.next());
                signatureGenerator.setHashedSubpackets(spGen.generate());
            }
            signatureGenerator.generateOnePassVersion(true).encode(compressedOut);

            // 字符数据生成器
            PGPLiteralDataGenerator literalDataGenerator = new PGPLiteralDataGenerator();
            // 包装成固定格式的字符输出流，写完就关
            //      1，借助缓冲数组写时，需要调用 iteralOut.close() 才能重写入数据长度（固定格式）
            OutputStream literalOut = literalDataGenerator.open(compressedOut, PGPLiteralData.BINARY, "pgp", new Date(), new byte[4096]);
            final byte[] buf = new byte[4096];
            for (int len; (len = in.read(buf)) >= 0;) {
                literalOut.write(buf, 0, len);
                signatureGenerator.update(buf, 0, len);
            }

            //      2，一开始就拿到了数据长度，正常写即可
            // OutputStream literalOut = literalDataGenerator.open(compressedOut, PGPLiteralData.BINARY, "pgp", clearData.length, new Date());
            // literalOut.write(clearData);
            // signatureGenerator.update(clearData);

            literalOut.close();

            signatureGenerator.generate().encode(compressedOut);
            compressedOut.close();
            encryptedOut.close();
            armored.close();

            return out.toString("UTF-8");
        }
    }

    /**
     * 解密  <p></p>
     *
     * @param encryptedStr  密文
     * @param filePath  私钥文件路径
     * @param password  密码
     * @return
     * @throws Exception
     */
    public static String decrypt(String encryptedStr, String filePath, String password) throws Exception {
        JcaPGPObjectFactory objectFactory =
                new JcaPGPObjectFactory(PGPUtil.getDecoderStream(new ByteArrayInputStream(encryptedStr.getBytes(StandardCharsets.UTF_8))));
        // 加密数据中的第一部分可能是 PGP标记包（可选的）
        Object obj = objectFactory.nextObject();
        PGPEncryptedDataList enc;
        if (obj instanceof PGPEncryptedDataList) {
            enc = (PGPEncryptedDataList) obj;
        } else {
            // 如果第一部分是 标记包，则获取第二部分数据
            enc = (PGPEncryptedDataList) objectFactory.nextObject();
        }

        // 获取 PGP加密数据
        PGPPublicKeyEncryptedData pbe = (PGPPublicKeyEncryptedData) enc.getEncryptedDataObjects().next();

        // 根据 密钥id 查找对应的私钥
        PGPPrivateKey privateKey = findPrivateKeyById(filePath, pbe.getKeyID(), password);

        // 通过私钥构建解密器
        BcPublicKeyDataDecryptorFactory decryptorFactory = new BcPublicKeyDataDecryptorFactory(privateKey);
        // 获得解密后数据输入流
        InputStream clear = pbe.getDataStream(decryptorFactory);
        // 通过解密后的输入流构建 PGP对象工厂
        objectFactory = new JcaPGPObjectFactory(clear);
        // 获取对象工厂第一部分的压缩数据
        PGPCompressedData message = (PGPCompressedData) objectFactory.nextObject();
        // 通过压缩数据构建 PGP对象工厂
        objectFactory = new JcaPGPObjectFactory(message.getDataStream());

        // 1，PGP对象的第一部分：签名方法
        PGPOnePassSignatureList algorithmList = (PGPOnePassSignatureList) objectFactory.nextObject();
        // 2，PGP对象的第二部分：消息明文
        PGPLiteralData pgpLiteralData = (PGPLiteralData) objectFactory.nextObject();

        // 解密后的明文
        InputStream input = pgpLiteralData.getInputStream();
        byte[] decrypted = Streams.readAll(input);


        return new String(decrypted, StandardCharsets.UTF_8);
    }

    /**
     * 解密并验证签名，因为验签需要源消息，需要先解密    <p></p>
     *
     * @param encryptedStr  密文
     * @param privateKey    私钥
     * @param publicKey     对方的公钥
     * @return
     * @throws Exception
     */
    public static String decryptAndVerify(String encryptedStr, PGPPrivateKey privateKey, PGPPublicKey publicKey) throws Exception {
        JcaPGPObjectFactory objectFactory =
                new JcaPGPObjectFactory(PGPUtil.getDecoderStream(new ByteArrayInputStream(encryptedStr.getBytes(StandardCharsets.UTF_8))));
        // 加密数据中的第一部分可能是 PGP标记包（可选的）
        Object obj = objectFactory.nextObject();
        PGPEncryptedDataList enc;
        if (obj instanceof PGPEncryptedDataList) {
            enc = (PGPEncryptedDataList) obj;
        } else {
            // 如果第一部分是 标记包，则获取第二部分数据
            enc = (PGPEncryptedDataList) objectFactory.nextObject();
        }

        // 获取 PGP加密数据
        PGPPublicKeyEncryptedData pbe = (PGPPublicKeyEncryptedData) enc.getEncryptedDataObjects().next();

        // 通过私钥构建解密器
        // PublicKeyDataDecryptorFactory decryptorFactory = new JcePublicKeyDataDecryptorFactoryBuilder().setProvider("BC").build(privateKey);
        BcPublicKeyDataDecryptorFactory decryptorFactory = new BcPublicKeyDataDecryptorFactory(privateKey);
        // 解密后数据输入流
        InputStream clear = pbe.getDataStream(decryptorFactory);
        // 通过解密后的输入流构建 PGP对象工厂
        objectFactory = new JcaPGPObjectFactory(clear);
        // 获取对象工厂第一部分的压缩数据
        PGPCompressedData message = (PGPCompressedData) objectFactory.nextObject();
        // 解压缩，并构建 PGP对象工厂
        objectFactory = new JcaPGPObjectFactory(message.getDataStream());

        // 1，PGP对象的第一部分：签名方法
        PGPOnePassSignatureList onePassList = (PGPOnePassSignatureList) objectFactory.nextObject();
        // 2，PGP对象的第二部分：消息明文
        PGPLiteralData pgpLiteralData = (PGPLiteralData) objectFactory.nextObject();

        // 解密后的明文
        InputStream input = pgpLiteralData.getInputStream();
        byte[] decrypted = Streams.readAll(input);

        // 获取签名对象，验证签名
        PGPOnePassSignature signature = onePassList.get(0);
        // 初始化公钥
        // signature.init(new JcaPGPContentVerifierBuilderProvider().setProvider("BC"), publicKey);
        signature.init(new BcPGPContentVerifierBuilderProvider(), publicKey);
        // 设置待验签数据（源数据）
        signature.update(decrypted);

        // 消息签名，只有读取消息体之后才能读到签名
        PGPSignatureList signatureList = (PGPSignatureList) objectFactory.nextObject();
        // 验证签名
        if (!signature.verify(signatureList.get(0))) {
            throw new Exception("pgp verify sign error.");
        }

        if (pbe.isIntegrityProtected() && !pbe.verify()) {
            LOG.error("Data is integrity protected but integrity is lost.");
        }
        return new String(decrypted, StandardCharsets.UTF_8);
    }

    /**
     * 从指定文件读取 完整的PGP密钥对    <p></p>
     *
     * @param filePath 私钥文件路径
     */
    public static PGPSecretKey readSecretKey(String filePath) throws Exception {
        InputStream input = new BufferedInputStream(Objects.requireNonNull(PGPUtils.class.getClassLoader().getResourceAsStream(filePath)));
        PGPSecretKeyRingCollection pgpSec = new PGPSecretKeyRingCollection(PGPUtil.getDecoderStream(input), new JcaKeyFingerprintCalculator());
        for (PGPSecretKeyRing keyRing : pgpSec) {
            for (PGPSecretKey key : keyRing) {
                if (key.isSigningKey()) {
                    return key;
                }
            }
        }

        throw new IllegalArgumentException("Can't find signing key in key ring.");
    }

    /**
     * 从指定文件读取 PGPSecretKey，并提取私钥    <p></p>
     *
     * @param filePath  私钥文件路径
     * @param password  密码
     */
    public static PGPPrivateKey findPrivateKey(String filePath, String password) throws Exception {
        PGPSecretKey pgpSecKey = readSecretKey(filePath);
        // 提取私钥
        return pgpSecKey.extractPrivateKey(new JcePBESecretKeyDecryptorBuilder().setProvider("BC").build(password.toCharArray()));
    }

    /**
     * 从指定文件读取 PGPSecretKey，并提取私钥    <p></p>
     *
     * @param pgpSecKey  完整的 PGP密钥对象
     * @param password  密码
     */
    public static PGPPrivateKey findPrivateKeyFromSecretKey(PGPSecretKey pgpSecKey, String password) throws Exception {
        // 提取私钥
        return pgpSecKey.extractPrivateKey(new JcePBESecretKeyDecryptorBuilder().setProvider("BC").build(password.toCharArray()));
    }

    /**
     * 从指定文件读取 完整的PGP密钥    <p></p>
     *
     * @param filePath 私钥文件路径
     * @param keyID     密钥id
     */
    public static PGPSecretKey readSecretKeyById(String filePath, long keyID) throws Exception {
        InputStream input = new BufferedInputStream(PGPUtils.class.getClassLoader().getResourceAsStream(filePath));
        PGPSecretKeyRingCollection pgpSec = new PGPSecretKeyRingCollection(PGPUtil.getDecoderStream(input), new JcaKeyFingerprintCalculator());
        PGPSecretKey secretKey = pgpSec.getSecretKey(keyID);
        if (secretKey == null) {
            throw new IllegalArgumentException("Can't find signing key in key ring.");
        }

        return secretKey;
    }

    /**
     * 从指定文件读取 PGPSecretKey，并提取私钥    <p></p>
     *
     * @param filePath  私钥文件路径
     * @param keyID     密钥id
     * @param password  密码
     */
    public static PGPPrivateKey findPrivateKeyById(String filePath, long keyID, String password) throws Exception {
        PGPSecretKey pgpSecKey = readSecretKeyById(filePath, keyID);
        // 提取私钥
        return pgpSecKey.extractPrivateKey(new JcePBESecretKeyDecryptorBuilder().setProvider("BC").build(password.toCharArray()));
    }

    /**
     * 从指定文件读取公钥
     */
    public static PGPPublicKey readPublicKey(String filePath) throws Exception {
        InputStream input = new BufferedInputStream(PGPUtils.class.getClassLoader().getResourceAsStream(filePath));
        PGPPublicKeyRingCollection pgpPub = new PGPPublicKeyRingCollection(PGPUtil.getDecoderStream(input), new JcaKeyFingerprintCalculator());

        for (PGPPublicKeyRing keyRing : pgpPub) {
            for (PGPPublicKey key : keyRing) {
                if (key.isEncryptionKey()) {
                    return key;
                }
            }
        }

        throw new IllegalArgumentException("Can't find encryption key in key ring.");
    }

    /**
     * 生成 PGP密钥对    <p></p>
     *
     * @param keyId                 即将生成的私钥编号
     * @param password              私钥密码
     * @param privateKeyFileName    私钥文件名
     * @param publicKeyFileName     公钥文件名
     */
    public static void generatePGPKeys(KeyId keyId, String password, String privateKeyFileName, String publicKeyFileName) throws Exception {
        final int s2kCount = 192;
        if (LOG.isDebugEnabled()) {
            LOG.debug("S2K Function Iteration Count : " + s2kCount);
        }

        Path path = Paths.get(Paths.get("").toAbsolutePath() + "/testJDK8/src/main/resources/cert/pgp");
        if (!Files.exists(path)) {
            Files.createDirectory(path);
        }

        // 初始化 PGP密钥对
        PGPKeyRingGenerator keyRingGenerator = generateKeyRingGenerator(keyId, password, s2kCount);

        // 保存公钥
        PGPPublicKeyRing pkr = keyRingGenerator.generatePublicKeyRing();
        ArmoredOutputStream pubOut = new ArmoredOutputStream(new BufferedOutputStream(
                Files.newOutputStream(Paths.get(path.toAbsolutePath() + File.separator + publicKeyFileName))));
        pkr.encode(pubOut, true);
        pubOut.close();

        String expireTimeLog;
        if (pkr.getPublicKey().getValidSeconds() != 0) {
            LocalDateTime dateTime = new Timestamp(pkr.getPublicKey().getCreationTime().getTime()).toLocalDateTime();
            dateTime = dateTime.plusSeconds(pkr.getPublicKey().getValidSeconds());
            expireTimeLog = "Generated PGP Public Key will expire on : "
                    + Date.from(dateTime.atZone(ZoneId.systemDefault()).toInstant());
        } else {
            expireTimeLog = "Generated PGP Public Key will expire on : Never";
        }
        if (LOG.isDebugEnabled()) {
            LOG.debug(expireTimeLog);
        }

        // 保存私钥
        PGPSecretKeyRing skr = keyRingGenerator.generateSecretKeyRing();
        ArmoredOutputStream secretOut = new ArmoredOutputStream(new BufferedOutputStream(
                Files.newOutputStream(Paths.get(path.toAbsolutePath() + File.separator + privateKeyFileName))));
        skr.encode(secretOut);
        secretOut.close();
        if (skr.getPublicKey().getValidSeconds() != 0) {
            LocalDateTime dateTime = new Timestamp(skr.getPublicKey().getCreationTime().getTime()).toLocalDateTime();
            dateTime = dateTime.plusSeconds(skr.getPublicKey().getValidSeconds());
            expireTimeLog = "Generated PGP Private Key will expire on : "
                    + Date.from(dateTime.atZone(ZoneId.systemDefault()).toInstant());
        } else {
            expireTimeLog = "Generated PGP Private Key will expire on : Never";
        }
        if (LOG.isDebugEnabled()) {
            LOG.debug(expireTimeLog);
        }
    }

    private static class KeyId {
        private final String keyId;

        public KeyId(String keyId) {
            this.keyId = keyId;
        }

        public String get() {
            return keyId;
        }
    }

    static {
        Security.addProvider(new BouncyCastleProvider());
    }

}
